package main

import (
	"fmt"
	"math"
	"math/rand"
)

const (
	maxLevel = 10
	p        = 0.5
)

type Node struct {
	Val     int
	Forward []*Node
}

func (l *Node) numLevel() int {
	return len(l.Forward)
}

type Skiplist struct {
	Head *Node
}

func Constructor() Skiplist {
	head := Node{
		Val:     math.MinInt32,
		Forward: make([]*Node, maxLevel),
	}
	return Skiplist{
		Head: &head,
	}
}

func (s *Skiplist) Search(target int) bool {
	cur := s.Head
	i := cur.numLevel() - 1
	t := 0
	for cur != nil && i >= 0 && cur.Val < target {
		if cur.Forward[i] == nil || cur.Forward[i].Val > target {
			i--
			t++
		} else {
			cur = cur.Forward[i]
			t++
		}
	}
	fmt.Println(t)
	return cur.Val == target
}

func (s *Skiplist) Add(num int) {
	cur := s.Head
	track := []*Node{cur}

	for level := cur.numLevel() - 1; level >= 0; {
		if cur.Forward[level] == nil || cur.Forward[level].Val >= num {
			level--
		} else {
			cur = cur.Forward[level]
			track = append(track, cur)
		}
	}

	// Head -> ... -> cur -> node -> ...
	node := &Node{
		Val:     num,
		Forward: make([]*Node, coinflipLayer()),
	}

	// Connect previous nodes to new node.
	curLevel := 0
	for i := len(track) - 1; i >= 0; i-- {
		prev := track[i]
		for curLevel < prev.numLevel() {
			node.Forward[curLevel] = prev.Forward[curLevel]
			prev.Forward[curLevel] = node
			curLevel++
			if curLevel >= node.numLevel() {
				return
			}
		}
	}

	// Connect new node to subsequent nodes.
	for next := node.Forward[curLevel]; next != nil && curLevel < node.numLevel(); {
		if curLevel < next.numLevel() {
			node.Forward[curLevel] = next
			curLevel++
		} else {
			next = next.Forward[curLevel]
		}
	}
}

func (s *Skiplist) Erase(num int) bool {
	// Each node's number of level in track will be in descending order.
	cur := s.Head
	track := []*Node{cur}

	// Go down & right to find the node with value = num.
	for level := len(cur.Forward) - 1; level >= 0; {
		if cur.Forward[level] == nil || cur.Forward[level].Val >= num {
			level--
		} else {
			cur = cur.Forward[level]
			track = append(track, cur)
		}
	}
	if cur.Forward[0] == nil || cur.Forward[0].Val != num {
		return false
	}
	rmNode := cur.Forward[0]
	for i, curLevel := len(track)-1, 0; i >= 0; i-- {
		prev := track[i]
		for curLevel < prev.numLevel() && curLevel < rmNode.numLevel() {
			prev.Forward[curLevel] = rmNode.Forward[curLevel]
			curLevel++
		}
	}
	return true
}

func coinflipLayer() int {
	res := 1
	for res < maxLevel-1 && rand.Float32() < p {
		res++
	}
	return res
}

func main() {
	l := Constructor()
	l.Add(1)
	l.Add(2)
	l.Add(3)
	l.Add(4)
	l.Add(5)
	l.Add(6)
	l.Add(7)
	l.Add(8)
	l.Add(9)
	l.Add(10)
	l.Add(11)
	l.Add(12)
	ret := l.Search(10)
	fmt.Println(ret)
}
